"""
会议调度环境管理器

集成到verl框架的环境管理接口
"""

import numpy as np
from typing import Dict, List, Any, Tuple

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

from agentic_system.environments.base import EnvironmentManagerBase
from .envs import MeetingSchedulingMultiProcessEnv
from .projection import meeting_projection, meeting_memory_projection


class MeetingSchedulingEnvManager(EnvironmentManagerBase):
    """
    会议调度环境管理器

    符合verl框架的EnvironmentManagerBase接口
    """

    def __init__(self,
                 envs: MeetingSchedulingMultiProcessEnv,
                 projection_f=None,
                 env_name: str = "meeting_scheduling"):
        """
        参数:
            envs: 矢量化环境实例
            projection_f: 投影函数（默认使用meeting_projection）
            env_name: 环境名称
        """
        # 使用默认投影函数
        if projection_f is None:
            projection_f = meeting_projection

        super().__init__(envs, projection_f, env_name)

        self.envs = envs
        self.projection_f = projection_f
        self.env_name = env_name

        # 环境状态
        self.current_text_observations = []
        self.episode_histories = []

    def reset(self) -> Dict[str, Any]:
        """
        重置环境

        返回:
            observations: 符合框架要求的观察字典
                {
                    'text': List[str],  # 文本观察
                    'image': np.ndarray,  # 向量观察（这里为None）
                    'anchor': None  # 原始观察
                }
        """
        # 重置所有环境
        text_observations, infos = self.envs.reset()

        self.current_text_observations = text_observations

        # 初始化episode histories
        self.episode_histories = [[] for _ in range(len(self.envs))]

        # 构建符合框架的观察格式
        observations = {
            'text': text_observations,
            'image': None,  # 会议调度主要使用文本，不需要向量观察
            'anchor': None
        }

        return observations

    def step(self, text_actions: List[str]) -> Tuple[Dict[str, Any], np.ndarray, np.ndarray, List[Dict]]:
        """
        执行步骤

        参数:
            text_actions: LLM生成的文本动作列表

        返回:
            (observations, rewards, dones, infos)
            - observations: 观察字典 {'text': List[str], 'image': None, 'anchor': None}
            - rewards: 奖励数组 shape=(env_num,)
            - dones: 完成标志数组 shape=(env_num,)
            - infos: 信息字典列表
        """
        # 投影文本动作
        action_dicts, valids = self.projection_f(text_actions)

        # 执行环境步骤
        text_observations, rewards_list, dones_list, infos = self.envs.step(
            text_actions,
            valids
        )

        self.current_text_observations = text_observations

        # 记录历史
        for i, info in enumerate(infos):
            if 'action' in info:
                self.episode_histories[i].append({
                    'step': info.get('step', 0),
                    'action_type': info.get('action'),
                    'observation': text_observations[i],
                    'reward': rewards_list[i]
                })

        # 转换为numpy数组
        rewards = np.array(rewards_list, dtype=np.float32)
        dones = np.array(dones_list, dtype=bool)

        # 构建观察
        observations = {
            'text': text_observations,
            'image': None,
            'anchor': None
        }

        return observations, rewards, dones, infos

    def success_evaluator(self, **kwargs) -> Dict[str, np.ndarray]:
        """
        成功率评估器

        根据info中的'won'标志评估成功率

        返回:
            {'success': np.ndarray}  # shape=(env_num,)
        """
        infos = kwargs.get('infos', [])

        success_flags = []
        for info in infos:
            # 检查是否有'won'字段
            won = info.get('won', False)
            success_flags.append(1 if won else 0)

        success = np.array(success_flags, dtype=np.int32)

        return {'success': success}

    def get_episode_histories(self) -> List[List[Dict]]:
        """获取所有环境的episode历史"""
        return self.episode_histories

    def get_current_observations(self) -> List[str]:
        """获取当前观察"""
        return self.current_text_observations


def create_meeting_env_manager(config: Dict) -> MeetingSchedulingEnvManager:
    """
    工厂函数：创建会议调度环境管理器

    参数:
        config: 配置字典，包含：
            - seed: 随机种子
            - env_num: 环境数量
            - max_steps: 最大步数
            - dataset_path: 数据集路径（可选）
            - is_train: 是否为训练模式
            - use_memory: 是否使用记忆投影函数

    返回:
        MeetingSchedulingEnvManager实例
    """
    seed = config.get('seed', 42)
    env_num = config.get('env_num', 4)
    max_steps = config.get('max_steps', 20)
    dataset_path = config.get('dataset_path', None)
    is_train = config.get('is_train', True)
    use_memory = config.get('use_memory', False)

    # 创建矢量化环境
    envs = MeetingSchedulingMultiProcessEnv(
        seed=seed,
        env_num=env_num,
        max_steps=max_steps,
        dataset_path=dataset_path,
        is_train=is_train
    )

    # 选择投影函数
    projection_f = meeting_memory_projection if use_memory else meeting_projection

    # 创建环境管理器
    env_manager = MeetingSchedulingEnvManager(
        envs=envs,
        projection_f=projection_f,
        env_name="meeting_scheduling"
    )

    return env_manager


def test_env_manager():
    """测试环境管理器"""
    print("=== 测试环境管理器 ===\n")

    # 配置
    config = {
        'seed': 42,
        'env_num': 2,
        'max_steps': 20,
        'is_train': True,
        'use_memory': False
    }

    # 创建环境管理器
    env_manager = create_meeting_env_manager(config)

    print(f"环境名称: {env_manager.env_name}")
    print(f"环境数量: {len(env_manager.envs)}\n")

    # 重置
    print("=== 重置环境 ===")
    observations = env_manager.reset()

    print(f"观察类型: {observations.keys()}")
    print(f"文本观察数量: {len(observations['text'])}")
    print(f"\n环境0观察: {observations['text'][0][:150]}...\n")

    # 步骤1
    print("=== 步骤1: 查询日程 ===")
    actions = [
        '''
        查询所有参会人员的日程。
        ```json
        {
            "action_type": "query_calendar",
            "parameters": {
                "persons": ["李明", "王芳", "陈静"]
            }
        }
        ```
        ''',
        '''
        查询必选人员的日程。
        ```json
        {
            "action_type": "query_calendar",
            "parameters": {
                "persons": ["李明", "王芳"]
            }
        }
        ```
        '''
    ]

    observations, rewards, dones, infos = env_manager.step(actions)

    print(f"奖励: {rewards}")
    print(f"完成标志: {dones}")
    print(f"\n环境0观察: {observations['text'][0][:200]}...\n")

    # 步骤2
    print("=== 步骤2: 查询历史 ===")
    actions = [
        '''
        ```json
        {
            "action_type": "query_history",
            "parameters": {}
        }
        ```
        ''',
        '''
        ```json
        {
            "action_type": "query_history",
            "parameters": {}
        }
        ```
        '''
    ]

    observations, rewards, dones, infos = env_manager.step(actions)

    print(f"奖励: {rewards}")
    print(f"完成标志: {dones}")

    # 成功率评估
    success_results = env_manager.success_evaluator(infos=infos)
    print(f"\n成功标志: {success_results['success']}")

    # 获取历史
    histories = env_manager.get_episode_histories()
    print(f"\n环境0历史步骤数: {len(histories[0])}")

    print("\n测试完成")


if __name__ == '__main__':
    test_env_manager()
